from tools.bus.vehicle import Vehicle
import cv2
import numpy as np

class VBus:
    def __init__(self,roi_list,iou_v,positive_vec) -> None:
        self.roi_list = np.array(roi_list)
        self.iou_v = iou_v
        self.positive_vec = positive_vec
        self.v_cnt = 1
        self.v_dict = {}
        
    def in_roi(self,box):
        return cv2.pointPolygonTest(self.roi_list, ((box[0]+box[2])/2,(box[1]+box[3])/2), False)>0
    
    def process(self,boxes,classes):
        caplist = [] #需要抓拍的车辆的(vtype,direction)
        num = len(boxes)
        for i in range(num):
            if not self.in_roi(boxes[i]):
                continue
            idx = 0
            iou = 0
            for key in self.v_dict.keys():
                t_iou = self.v_dict[key].get_iou(boxes[i],classes[i])
                if t_iou>self.iou_v and t_iou>iou:
                    idx=key
                    iou=t_iou
            if idx > 0:
                #已经检测的车
                if not self.v_dict[idx].captured:
                    self.v_dict[idx].update(boxes[i],self.positive_vec)
                    caplist.append((classes[i],self.v_dict[idx].direction))
                else:
                    self.v_dict[idx].update(boxes[i],self.positive_vec)
            else:
                # 新车
                self.v_dict[self.v_cnt]=Vehicle(boxes[i],classes[i],self.v_cnt)
                self.v_cnt+=1
        self.update()
        return caplist
    
    def update(self):
        del_keys=[]
        for key in self.v_dict.keys():
            if not self.v_dict[key].updated:
                del_keys.append(key)
                continue
            else:
                self.v_dict[key].updated=False    
        for key in del_keys:
            del self.v_dict[key]   
            